#!/usr/bin/env python
# $Id: ssh.py 6525fa8d3382 2012/09/10 23:32:19 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

__version__  = "$Revision: 6525fa8d3382 $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2012-09-10 23:32:19Z $"

# system imports
import getpass
import os
import optparse
import shlex
import signal
import sys

# local imports
# import cmdparser library to handle command line arguments
from venya.lib import cmdparser

# Import the core of the ssh implementation using paramiko
from venya.lib import sshlib
from venya.lib import log
from venya.lib import stdlib
from venya.lib import console

# Main client class, which imports the command parser class, the actual
# ssh handling is done using the handle function

# Main SSH client class
class sshclient(cmdparser.GetOpt):
   option_list = (
         optparse.make_option("-4", action = "store_true",
            help = "Use IPv4 for connecting", default = False, dest = "v4"),
         optparse.make_option("-6", action = "store_true",
            help = "Use IPv6 for connecting", default = False, dest = "v6"),
         optparse.make_option("-l", "--login", type = "str",
            help = "Username for SSH session", dest = "login"),
         optparse.make_option("-L", "--list", type = "str",
            help = "List of servers to connect", dest = "syslist"),
         optparse.make_option("-f", "--ftp", action = "store_true",
            help = "SFTP mode using SSH", default = False,
            dest = "ftp"),
         optparse.make_option("-F", "--filename", type = "str",
            help = "Filename to parse and send to SSH server",
            dest = "filename"),
         optparse.make_option("-p", "--port", type = "int",
            help = "SSH port to use, default: 22", default = 22,
            dest = "port"),
         optparse.make_option("-P", "--password", action = "store_true",
            help = "Prompt for password", dest = "password"),
         ) + cmdparser.GetOpt.option_list

   help = """
   SSH and SFTP client using paramiko
   """

   def __init__(self, **keywords):
      super(self.__class__, self).__init__(**keywords)
      self._log = log.StreamLog(stdlib.my_name(self))
      self._client = None
      self._host = None
      self._cmdline = None
      self._cmdlen = 0
      completer = console.CmdCompleter({
         'get' : [],
         'lcd': [],
         'put': [],
         'exit' : [],
         'lpwd': [],
         'lls': [],
         'stat': [],
         })
      self._console = console.CmdLine(completer)

   def sighandle(self, signal, frame):
      print("\n")
      self.error("program with pid %d terminated with signal: %d",
            os.getpid(), signal)

   def setsignal(self, flag = True):
      func = None
      if flag:
         func = self.sighandle
      else:
         func = signal.SIG_IGN

      try:
         signal.signal(signal.SIGTERM, func)
         signal.signal(signal.SIGINT, func)
         signal.signal(signal.SIGHUP, func)
      except ValueError, message:
         pass

   def _sftpget(self):
      try:
         if self._cmdlen == 1:
            raise AttributeError, "need more arguments"
         elif self._cmdlen == 2:
            local = self._cmdline[1].split("/")[-1]
            self._client.Get(self._cmdline[1], local)
         elif self._cmdlen == 3:
            self._client.Get(self._cmdline[1], self._cmdline[2])
         elif self._cmdlen > 3:
            raise AttributeError, "too many arguments"
      except IOError, (error, strerror):
         raise AttributeError, "%s: %s(%d)" %(self._cmdline[1], strerror, error)

   def _sftpput(self):
      print(self._cmdline)

   def _sftplcd(self):
      if self._cmdline:
         if self._cmdlen == 2:
            try:
               os.chdir(self._cmdline[1])
            except OSError, (error, strerror):
               raise AttributeError, "%s: %s, %s" %(
                     str(self), self._cmdline[1], strerror)
         elif self._cmdlen < 2:
            raise AttributeError, "need more arguments"
         elif self._cmdlen > 2:
            raise AttributeError, "%s: too many arguments"

   def sftp(self):
      self.setsignal(False)
      self._client.sftp()
      flag = True
      while flag:
         try:
            prompt = "%s [%d]> " %(self._host, self._console.GetHistLength())
            line = self._console.raw_input(prompt)
         except EOFError:
            flag = False
            print
         except KeyboardInterrupt:
            print
            pass
         else:
            if line:
               command = line.strip()
               self._cmdline = shlex.split(command)
               self._cmdlen = len(self._cmdline)
               cmd = self._cmdline[0].lower()
               try:
                  if cmd == "exit":
                     flag = False
                  elif cmd == "lcd":
                     self._sftplcd()
                  elif cmd == "get":
                     self._sftpget()
                  elif cmd == "put":
                     self._sftpput()
                  elif cmd == "stat":
                     self._sftpstat()
               except AttributeError, message:
                  print("%s: %s" %(str(self), message))

   @cmdparser.getopts
   def handle(self):
      username = None
      password = None
      filename = None
      port = 22
      servers = []

      self.setsignal()

      if self.options.syslist:
         try:
            with open(self.options.syslist) as fh:
               for lines in fh.readlines():
                  servers.append(lines.strip())
         except IOError as err:
            self.error(message)
      else:
         try:
            servers = self.arguments[1:]
         except IndexError, message:
            pass

      if len(servers) == 0:
         self.error("Cannot proceed without any host name!")

      if self.options.login:
         username = self.options.login.strip()
      else:
         username = getpass.getuser()

      if self.options.port:
         try:
            port = int(self.options.port)
         except ValueError, message:
            self.error(message)

      if self.options.password:
         if len(servers) == 1:
            password = getpass.getpass("%s password on %s: " %(
               username, servers[0]))
         else:
            password = getpass.getpass("%s password: " %username)

         password = password.strip()

      if self.options.filename:
         filename = self.options.filename.strip()

      for device in servers:
         try:
            with sshlib.SSHClient(host = device, user = username, port = port,
                  verbose = bool(self.options.verbosity),
                  passwd = password) as self._client:
               if self.options.v4:
                  self._client.Connect(family = 4)
               elif self.options.v6:
                  self._client.Connect(family = 6)
               else:
                  self._client.Connect()
   
               if self._client.Auth(bool(password)):
                  if self.options.ftp:
                     self.sftp()
                  else:
                     if not self.options.filename:
                        self.setsignal(False)

                     self._client.GetChannel(exfile = filename)
               else:
                  self.warn("cannot authenticate to %s", self._host)
         except sshlib.SSHException, message:
            self.warn(message)
         finally:
            self._client = None

if __name__ == "__main__":
   client = None
   try:
      client = sshclient(version = 1.4, args = sys.argv)
      if len(client) <= 1:
         client.printhelp()
      else:
         client.handle()
   except cmdparser.GetOptError, message:
      print(message)
      sys.exit(1)
